"""
动作生成模块
根据音频特征生成机器人动作序列
"""

import time
import numpy as np
from typing import Dict, List, Tuple, Any
from dataclasses import dataclass
from enum import Enum


class DanceMode(Enum):
    """舞蹈模式枚举"""
    PRESET = "preset"       # 预设模式
    REALTIME = "realtime"   # 实时模式
    RECORDING = "recording" # 录制模式
    EDITING = "editing"     # 编辑模式


@dataclass
class RobotAction:
    """机器人动作数据类"""
    timestamp: float
    joint_angles: Dict[str, float]  # 关节角度 {关节名: 角度值}
    duration: float  # 动作持续时间
    action_id: str = ""  # 动作ID


@dataclass
class DanceSequence:
    """舞蹈序列数据类"""
    actions: List[RobotAction]
    sequence_id: str = ""
    metadata: Dict[str, Any] = None


class MotionGenerator:
    """
    动作生成器
    负责根据音频特征生成机器人动作
    """
    
    def __init__(self):
        """初始化动作生成器"""
        self.dance_sequences = {}  # 存储预设舞蹈序列
        self.current_sequence = None
        self.current_action_index = 0
        self.mode = DanceMode.REALTIME
        
        # 初始化一些预设动作
        self._initialize_preset_actions()
        
    def _initialize_preset_actions(self):
        """初始化预设动作"""
        # 这里可以加载一些预设的舞蹈动作序列
        pass
        
    def generate_realtime_action(self, audio_features: Dict[str, Any]) -> RobotAction:
        """
        根据实时音频特征生成动作
        
        Args:
            audio_features: 音频特征数据
            
        Returns:
            RobotAction: 生成的机器人动作
        """
        # 获取当前时间戳
        timestamp = time.time()
        
        # 根据音频特征生成关节角度
        joint_angles = self._map_audio_to_joints(audio_features)
        
        # 根据节拍信息确定动作持续时间
        duration = self._calculate_action_duration(audio_features)
        
        return RobotAction(
            timestamp=timestamp,
            joint_angles=joint_angles,
            duration=duration,
            action_id=f"realtime_{int(timestamp * 1000)}"
        )
        
    def _map_audio_to_joints(self, audio_features: Dict[str, Any]) -> Dict[str, float]:
        """
        将音频特征映射到关节角度
        
        Args:
            audio_features: 音频特征数据
            
        Returns:
            Dict[str, float]: 关节角度字典
        """
        # 简单示例：根据音量和频率控制不同关节
        volume = audio_features.get('amplitude_rms', 0)
        frequency = audio_features.get('dominant_frequency', 0)
        is_beat = audio_features.get('is_beat', False)
        
        # 对音量和频率进行归一化处理，确保在合理范围内
        # 音量范围通常在0-100之间
        volume = max(0, min(100, volume))
        # 频率范围通常在0-1000Hz之间
        frequency = max(0, min(1000, frequency))
        
        # 根据音频特征生成关节角度
        joint_angles = {
            'left_arm': self._map_value(volume, 0, 100, -45, 45),
            'right_arm': self._map_value(frequency, 0, 1000, -45, 45),
            'left_leg': self._map_value(volume, 0, 100, -15, 15) if is_beat else 0,
            'right_leg': self._map_value(frequency, 0, 1000, -15, 15) if is_beat else 0,
            'head': self._map_value(frequency, 0, 1000, -45, 45)
        }
        
        return joint_angles
        
    def _calculate_action_duration(self, audio_features: Dict[str, Any]) -> float:
        """
        根据音频特征计算动作持续时间
        
        Args:
            audio_features: 音频特征数据
            
        Returns:
            float: 动作持续时间（秒）
        """
        # 简单示例：根据BPM计算动作持续时间
        bpm = audio_features.get('bpm', 120)
        if bpm and bpm > 0:
            # 每拍的时长（秒）
            beat_duration = 60.0 / bpm
            return max(0.1, beat_duration / 2)  # 动作持续半拍，但至少0.1秒
        else:
            return 0.5  # 默认0.5秒
            
    def _map_value(self, value: float, in_min: float, in_max: float, 
                   out_min: float, out_max: float) -> float:
        """
        将值从一个范围映射到另一个范围
        
        Args:
            value: 输入值
            in_min: 输入最小值
            in_max: 输入最大值
            out_min: 输出最小值
            out_max: 输出最大值
            
        Returns:
            float: 映射后的值
        """
        # 确保输入值在范围内
        value = max(in_min, min(in_max, value))
        
        # 线性映射
        return (value - in_min) * (out_max - out_min) / (in_max - in_min) + out_min
        
    def load_preset_sequence(self, sequence_id: str) -> DanceSequence:
        """
        加载预设舞蹈序列
        
        Args:
            sequence_id: 序列ID
            
        Returns:
            DanceSequence: 舞蹈序列
        """
        # 这里应该从文件或数据库加载预设序列
        # 现在返回一个示例序列
        actions = [
            RobotAction(
                timestamp=time.time(),
                joint_angles={'left_arm': -30, 'right_arm': 30, 'left_leg': 0, 'right_leg': 0, 'head': 0},
                duration=0.5,
                action_id="preset_1"
            ),
            RobotAction(
                timestamp=time.time() + 0.5,
                joint_angles={'left_arm': 30, 'right_arm': -30, 'left_leg': 0, 'right_leg': 0, 'head': 0},
                duration=0.5,
                action_id="preset_2"
            )
        ]
        
        sequence = DanceSequence(
            actions=actions,
            sequence_id=sequence_id,
            metadata={'name': '示例舞蹈', 'duration': 1.0}
        )
        
        self.dance_sequences[sequence_id] = sequence
        return sequence
        
    def get_next_preset_action(self) -> RobotAction:
        """
        获取下一个预设动作
        
        Returns:
            RobotAction: 下一个动作
        """
        if not self.current_sequence or not self.current_sequence.actions:
            return None
            
        if self.current_action_index < len(self.current_sequence.actions):
            action = self.current_sequence.actions[self.current_action_index]
            self.current_action_index += 1
            return action
        else:
            # 序列结束，重置索引
            self.current_action_index = 0
            return None
            
    def set_editing_sequence(self, actions_data: List[Dict]):
        """
        设置编辑模式下的动作序列
        
        Args:
            actions_data: 动作数据列表
        """
        print(f"设置编辑模式序列，动作数量: {len(actions_data)}")
        actions = []
        for i, data in enumerate(actions_data):
            action = RobotAction(
                timestamp=data.get('timestamp', i * 0.5),
                joint_angles=data.get('joint_angles', {}),
                duration=data.get('duration', 0.5),
                action_id=data.get('action_id', f"editing_{i}")
            )
            actions.append(action)
            print(f"添加动作 {i}: {action}")
            
        sequence = DanceSequence(
            actions=actions,
            sequence_id="editing_sequence",
            metadata={'name': '编辑模式序列', 'duration': len(actions) * 0.5}
        )
        
        self.dance_sequences["editing"] = sequence
        self.current_sequence = sequence
        self.current_action_index = 0
        print("编辑模式序列设置完成")
        
    def get_next_editing_action(self) -> RobotAction:
        """
        获取下一个编辑模式动作
        
        Returns:
            RobotAction: 下一个动作
        """
        if not self.current_sequence or not self.current_sequence.actions:
            print("编辑模式: 当前序列为空")
            return None
            
        if self.current_action_index < len(self.current_sequence.actions):
            action = self.current_sequence.actions[self.current_action_index]
            print(f"编辑模式: 获取动作 {self.current_action_index}: {action.action_id}")
            self.current_action_index += 1
            return action
        else:
            # 序列结束，返回None表示序列已完成
            print("编辑模式: 序列结束")
            return None


# 全局动作生成器实例
motion_generator = MotionGenerator()


def get_motion_generator() -> MotionGenerator:
    """
    获取全局动作生成器实例
    
    Returns:
        MotionGenerator: 动作生成器实例
    """
    return motion_generator